################################################################################
class rap_trans:
  """ 
  Class for storing transition data; also serves as a superclass for
  specific transition data, e.g. A-values.
  """
  def __init__(self,levs=None):
    self.clear()
    if levs is not None: self.set_levs(levs)

# ------------------------------------------------------------------------------

  def __len__(self):
    return self.num_transitions()

# ------------------------------------------------------------------------------

  def __getitem__(self,key):
    return self.v_data.__getitem__(key)

# ------------------------------------------------------------------------------

  def __contains__(self,key):
    return self.v_data.__contains__(key)

# ------------------------------------------------------------------------------

  def __iter__(self):
    return self.v_data.__iter__()

# ------------------------------------------------------------------------------

  def _relimit(self):
    """ Remove transitions not in given initial/final state limits """
    if self.v_iminlev is not None:
      for (i,j) in self.v_data.keys():
        if i < self.v_iminlev: del self.v_data[(i,j)]
    if self.v_imaxlev is not None:
      for (i,j) in self.v_data.keys():
        if i > self.v_imaxlev: del self.v_data[(i,j)]
    if self.v_fminlev is not None:
      for (i,j) in self.v_data.keys():
        if j < self.v_fminlev: del self.v_data[(i,j)]
    if self.v_fmaxlev is not None:
      for (i,j) in self.v_data.keys():
        if j > self.v_fmaxlev: del self.v_data[(i,j)]

# ------------------------------------------------------------------------------

  def clear(self,keeplevs=False):
    """ Reset data. """
    if not keeplevs: self.v_levs=None
    self.v_iminlev=None     # lower limit for initial state index
    self.v_imaxlev=None     # upper limit for initial state index
    self.v_fminlev=None     # lower limit for final state index
    self.v_fmaxlev=None     # upper limit for final state index
    self.v_data={}

# ------------------------------------------------------------------------------

  def reindex(self,remap):
    """ 
    Reassign level indices.  Indices are remapped according to dictionary,
    remap, old (key) -> new (value).  
    """
    indices=[]
    for (i,j) in self:
      if i not in indices: indices.append(i)
      if j not in indices: indices.append(j)
    err=False
    for i in indices:
      if i not in remap:
        err=True
        break
    if err:
      raise ErrorRAP("all indices must be represented using reindex")
    tdat={}
    for (i,j),val in self.v_data.iteritems(): tdat[(remap[i],remap[j])]=val
    self.v_data=tdat
    if self.v_iminlev is not None: self.v_iminlev=remap[self.v_iminlev]
    if self.v_imaxlev is not None: self.v_imaxlev=remap[self.v_imaxlev]
    if self.v_fminlev is not None: self.v_fminlev=remap[self.v_fminlev]
    if self.v_fmaxlev is not None: self.v_fmaxlev=remap[self.v_fmaxlev]
    return remap

# ------------------------------------------------------------------------------

  def set_levs(self,levs):
    """ Set rap_states object. """
    if not isinstance(levs,rap_states):
      msg='levs must be a rap_states instance (or child instance)'
      raise ErrorRAP(msg)
    self.v_levs=levs

# ------------------------------------------------------------------------------

  def levs(self):
    """ Return rap_states object. """
    return self.v_levs

# ------------------------------------------------------------------------------

  def num_levels(self):
    """ Return number of levels. """
    return len(self.v_levs)

# ------------------------------------------------------------------------------

  def num_transitions(self):
    """ Return number of transitions. """
    return len(self.v_data)

# ------------------------------------------------------------------------------

  def initial_limit(self,imap=False):
    """ Return allowed range for inital state indices; None for no limit. """
    imin=self.v_iminlev
    imax=self.v_imaxlev
    if imap is False: return (imin,imax)
    if imin is None and imax is None: return (imin,imax)
    tmin=1
    if imin is not None: tmin=imin
    tmax=self.num_levels()
    if tmax is not None: tmax=imax
    levlist=[self.v_levs.lmap_reverse(i) for i in range(tmin,tmax+1)]
    tmin=min(levlist)
    tmax=max(levlist)
    if imin is None: tmin=None
    if imax is None: tmax=None
    return (tmin,tmax)

# ------------------------------------------------------------------------------

  def final_limit(self,imap=False):
    """ Return allowed range for final state indices; None for no limit. """
    fmin=self.v_fminlev
    fmax=self.v_fmaxlev
    if imap is False: return (fmin,fmax)
    if fmin is None and fmax is None: return (fmin,fmax)
    tmin=1
    if fmin is not None: tmin=fmin
    tmax=self.num_levels()
    if tmax is not None: tmax=fmax
    levlist=[self.v_levs.lmap_reverse(i) for i in range(tmin,tmax+1)]
    tmin=min(levlist)
    tmax=max(levlist)
    if fmin is None: tmin=None
    if fmax is None: tmax=None
    return (tmin,tmax)

# ------------------------------------------------------------------------------

  def set_initial_limit(self,minlev=None,maxlev=None,imap=False):
    """ Set allowed range for inital state indices; None for no limit. """
    tmin=None
    if minlev is not None: tmin=self.v_levs._mapme(minlev,imap)
    tmax=None
    if maxlev is not None: tmax=self.v_levs._mapme(maxlev,imap)
    if tmin is not None and tmax is not None:
      if tmin > tmax: tmin,tmax=tmax,tmin
    indices=self.v_levs.indices()
    if tmin is not None and tmin < min(indices):
      msg='minlev out of range'
      raise ErrorRAP(msg)
    if tmax is not None and tmax > max(indices):
      msg='maxlev out of range'
      raise ErrorRAP(msg)
    self.v_iminlev=tmin
    self.v_imaxlev=tmax
    self._relimit()

# ------------------------------------------------------------------------------

  def set_final_limit(self,minlev=None,maxlev=None,imap=False):
    """ Set allowed range for final state indices; None for no limit. """
    tmin=None
    if minlev is not None: tmin=self.v_levs._mapme(minlev,imap)
    tmax=None
    if maxlev is not None: tmax=self.v_levs._mapme(maxlev,imap)
    if tmin is not None and tmax is not None:
      if tmin > tmax: tmin,tmax=tmax,tmin
    indices=self.v_levs.indices()
    if tmin is not None and tmin < min(indices):
      msg='minlev out of range'
      raise ErrorRAP(msg)
    if tmax is not None and tmax > max(indices):
      msg='maxlev out of range'
      raise ErrorRAP(msg)
    self.v_fminlev=tmin
    self.v_fmaxlev=tmax
    self._relimit()

# ------------------------------------------------------------------------------

  def initial_indices(self,imap=False):
    """ Return list of active indices. """
    out=list(set([sublist[0] for sublist in self.v_data.keys()]))
    (imin,imax)=self.initial_limit()
    if imin is not None: out=filter(lambda x: x >= imin,out)
    if imax is not None: out=filter(lambda x: x <= imax,out)
    if imap: out=self.v_levs.lmap_reverse(out)
    return sorted(out)

# ------------------------------------------------------------------------------

  def final_indices(self,imap=False):
    """ Return list of active indices. """
    out=list(set([sublist[1] for sublist in self.v_data.keys()]))
    (imin,imax)=self.initial_limit()
    if imin is not None: out=filter(lambda x: x >= imin,out)
    if imax is not None: out=filter(lambda x: x <= imax,out)
    if imap: out=self.v_levs.lmap_reverse(out)
    return sorted(out)

# ------------------------------------------------------------------------------

  def indices(self,imap=False):
    """ Return list of active indices. """
    out=self.initial_indices()
    out.extend(self.final_indices())
    out=list(set(out))
    if imap: out=self.v_levs.lmap_reverse(out)
    return sorted(out)

# ------------------------------------------------------------------------------

  def initial_filter(self,istates,imap=False):
    """ Return list of active indices. """
    tmp=istates
    if not isinstance(tmp,list): tmp=[tmp]
    out=[key for key in self.v_data.keys() if key[0] in tmp]
    if imap: out=self.v_levs.lmap_reverse(out)
    return sorted(out)

# ------------------------------------------------------------------------------

  def final_filter(self,istates,imap=False):
    """ Return list of active indices. """
    tmp=istates
    if not isinstance(tmp,list): tmp=[tmp]
    out=[key for key in self.v_data.keys() if key[1] in tmp]
    if imap: out=self.v_levs.lmap_reverse(out)
    return sorted(out)

# ------------------------------------------------------------------------------

  def transitions(self,imap=False):
    """ Return list of transition tuples: (lower,upper). """
    out=self.v_data.keys()
    if imap: out=self.v_levs.lmap_reverse(out)
    return sorted(out)

# ------------------------------------------------------------------------------

  def data(self,i,j=None,imap=False):
    """ 
    Return the data value for the i-j transition.
    
    Note: it is assumed that state i has a lower energy than state j.
    """
    ii,jj=i,j
    if jj is None and isinstance(ii,tuple): (ii,jj)=ii
    trans=(ii,jj)
    if imap:
      trans=self.v_levs.lmap(trans)
      if trans is None: raise ErrorRAP('invalid mapping')
    if trans not in self.v_data: 
      (ii,jj)=trans
      raise ErrorRAP('Data value for %i-%i transition not found.'%(ii,jj))
    return self.v_data[trans]

# ------------------------------------------------------------------------------

  def set_data(self,i,j=None,value=None,imap=False):
    """ 
    Set data value for the i-j transition.
    
    Note: it is assumed that state i has a lower energy than state j.
    """
    if isinstance(i,tuple):
      (ii,jj)=i
      tval=j
    else:
      ii,jj=i,j
      tval=value
    if tval is None:
      raise ErrorRAP("need to give a new value")
    if ii not in self.levs():
      raise ErrorRAP("Invalid state index given: %s"%ii)
    if jj not in self.levs():
      raise ErrorRAP("Invalid state index given: %s"%jj)
    trans=(ii,jj)
    if imap:
      trans=self.v_levs.lmap(trans)
      if trans is None: raise ErrorRAP('invalid mapping')
    self.v_data[trans]=tval

# ------------------------------------------------------------------------------

  def wavelength(self,i,j,imap=False):
    """ Return wavelength calculated from level structure in Angstroms. """
    return self.levs().wavelength(i,j,imap=imap)

################################################################################

